# 绘制过拟合

import json
import configparser

import matplotlib.pyplot as plt
import os
import src.tool as tool
import numpy as np

# 配置文件地址
filename = './config.ini'
# 读取程序配置
config = configparser.ConfigParser()
config.read(filename, encoding='utf-8')

LR = tool.get_config_by_key(config['ai'], 'LR')
BATCH_SIZE = tool.get_config_by_key(config['ai'], 'BATCH_SIZE')

filename = './saved_model/overfitting.txt'

trainfilename = './saved_model/train_data.txt'


def paint(all_data: list, train_data: list):
    epoch_winner = []
    loss = []
    accuracy = []
    epoch = 0
    loss_mean = []
    train_loss = []
    for i in range(len(all_data)):
        single = all_data[i]
        try:
            parsed = json.loads(single)
        except Exception as e:
            print("第{}行不是需要的数据，错误原因->{}".format(i + 1, e))
            continue
        data_type = parsed['type']
        data = parsed['data']
        if data_type == 'loss':
            loss_mean.append(np.mean(data))
            epoch = epoch + 1
    for i in range(len(train_data)):
        single = train_data[i]
        try:
            parsed = json.loads(single)
        except Exception as e:
            print("第{}行不是需要的数据，错误原因->{}".format(i + 1, e))
            continue
        data_type = parsed['type']
        data = parsed['data']
        if data_type == 'loss':
            train_loss.extend(data)
            train_loss.append(np.mean(data))
            epoch = epoch + 1

    train_loss = train_loss[len(train_loss)-10:len(train_loss)]
    print("test_loss->",loss_mean)
    print("train_loss->", train_loss)
    fig, ax = plt.subplots()
    ax.plot(loss_mean, label='test-loss')
    ax.plot(train_loss, label='train-loss')
    ax.set_title('check over fit')
    ax.set_xlabel('LR={},batch_size = {},epoch = {}'.format(LR, BATCH_SIZE, epoch))
    ax.legend()
    plt.grid(True)
    plt.show()


if __name__ == '__main__':
    fs = open(filename, encoding='utf-8', mode='r')
    all_data = fs.read()
    lines = all_data.split('\n')
    fs_train = open(trainfilename, encoding='utf-8', mode='r')
    train_data = fs_train.read()
    train_lines = train_data.split('\n')
    paint(lines, train_lines)
